import Model.AutoEncoding as MyModel # 模型
import Loader as MyLoader # 数据加载器
import Config # 训练参数
from piqa import PSNR # 评价图像相似程度，越大越相似

import torch
from torch import nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


state_dict = torch.load(Config.BEST_MODEL_PATH)
model = MyModel.DenoiseAutoEncoder().to(device)
model.load_state_dict(state_dict['model'])
psnr = PSNR(reduction='sum') # 有三种，none，sum，mean，默认mean。

test_data_loader = MyLoader.loadTestData()

criterion = nn.MSELoss().to(device)


model = model.eval()
# 测试
batch_num = 0
total_loss = 0
total_psnr = 0
for batch_data in test_data_loader:
    X = batch_data['X']
    y = batch_data['y']
    outputs = model(X)

    outputs = model(X)
    loss = criterion(outputs, y)

    total_loss += loss.item() * X.size(0)
    batch_num += X.size(0)
    total_psnr += psnr(X,y)

print(f"测试的Ave MSELoss为{total_loss/batch_num}")
print(f"测试的Ave PSNR为{total_psnr/batch_num}")
        
